import numpy as np
import matplotlib.pyplot as plt

# 设置matplotlib支持中文显示 - 使用Windows系统常见中文字体
plt.rcParams["font.family"] = ["Microsoft YaHei", "SimHei", "SimSun", "FangSong", "KaiTi"]
plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号

class InfiniteSquareWell:
    """
    一维无限深势阱类，用于描述和计算量子力学中的一维无限深势阱系统
    
    属性:
        length: 势阱宽度L
        mass: 粒子质量m
        hbar: 约化普朗克常数
    """
    
    def __init__(self, length=1.0, mass=1.0, hbar=1.0):
        """
        初始化一维无限深势阱实例
        
        参数:
            length: 势阱宽度，默认为1.0
            mass: 粒子质量，默认为1.0
            hbar: 约化普朗克常数，默认为1.0
        """
        self.length = length  # 势阱宽度 L
        self.mass = mass      # 粒子质量 m
        self.hbar = hbar      # 约化普朗克常数 ħ
    
    def wavefunction(self, x, n):
        """
        计算一维无限深势阱中第n个能级的波函数
        
        波函数公式：ψ_n(x) = √(2/L) * sin(nπx/L)
        其中x ∈ [0, L]，L为势阱宽度，n为量子数(n=1,2,3,...)
        
        参数:
            x: 位置坐标或坐标数组
            n: 量子数
        
        返回:
            波函数在给定位置的值
        """
        # 确保x在势阱范围内[0, L]
        x = np.asarray(x)
        result = np.zeros_like(x, dtype=complex)
        
        # 在势阱内部计算波函数
        inside_well = (x >= 0) & (x <= self.length)
        result[inside_well] = np.sqrt(2 / self.length) * np.sin(n * np.pi * x[inside_well] / self.length)
        
        # 势阱外部波函数为0
        return result
    
    def energy(self, n):
        """
        计算一维无限深势阱中第n个能级的能量
        
        能量公式：E_n = (n²π²ħ²)/(2mL²)
        
        参数:
            n: 量子数
        
        返回:
            能量值
        """
        return (n**2 * np.pi**2 * self.hbar**2) / (2 * self.mass * self.length**2)
    
    def probability_density(self, x, n):
        """
        计算波函数的概率密度
        
        概率密度公式：|ψ_n(x)|²
        
        参数:
            x: 位置坐标或坐标数组
            n: 量子数
        
        返回:
            概率密度在给定位置的值
        """
        psi = self.wavefunction(x, n)
        return np.abs(psi)**2
    
    def plot_wavefunction(self, n, num_points=1000):
        """
        绘制第n个能级的波函数
        
        参数:
            n: 量子数
            num_points: 绘图的点数，默认为1000
        """
        x = np.linspace(-0.2 * self.length, 1.2 * self.length, num_points)
        psi = self.wavefunction(x, n)
        
        plt.figure(figsize=(10, 6))
        plt.plot(x, np.real(psi), 'b-', label=f'波函数 ψ_{n}(x)')
        plt.axvline(x=0, color='k', linestyle='--', alpha=0.5)
        plt.axvline(x=self.length, color='k', linestyle='--', alpha=0.5)
        plt.fill_betweenx([min(np.real(psi)), max(np.real(psi))], 0, self.length, color='gray', alpha=0.1, label='势阱区域')
        plt.title(f'一维无限深势阱的波函数 (n={n})')
        plt.xlabel('位置 x')
        plt.ylabel('波函数 ψ(x)')
        plt.grid(True)
        plt.legend()
        plt.show(block=True)
    
    def plot_probability_density(self, n, num_points=1000):
        """
        绘制第n个能级的概率密度
        
        参数:
            n: 量子数
            num_points: 绘图的点数，默认为1000
        """
        x = np.linspace(-0.2 * self.length, 1.2 * self.length, num_points)
        prob_density = self.probability_density(x, n)
        
        plt.figure(figsize=(10, 6))
        plt.plot(x, prob_density, 'r-', label=f'概率密度 |ψ_{n}(x)|²')
        plt.axvline(x=0, color='k', linestyle='--', alpha=0.5)
        plt.axvline(x=self.length, color='k', linestyle='--', alpha=0.5)
        plt.fill_betweenx([0, max(prob_density)], 0, self.length, color='gray', alpha=0.1, label='势阱区域')
        plt.title(f'一维无限深势阱的概率密度 (n={n})')
        plt.xlabel('位置 x')
        plt.ylabel('概率密度 |ψ(x)|²')
        plt.grid(True)
        plt.legend()
        plt.show(block=True)
    
    def plot_energy_levels(self, max_n=5):
        """
        绘制前max_n个能级的能量图
        
        参数:
            max_n: 要绘制的最高量子数，默认为5
        """
        plt.figure(figsize=(8, 10))
        
        for n in range(1, max_n + 1):
            energy = self.energy(n)
            plt.axhline(y=energy, xmin=0.2, xmax=0.8, color='blue', linewidth=2)
            plt.text(0.1, energy, f'E_{n} = {energy:.4f}', va='center', fontsize=12)
        
        plt.title('一维无限深势阱的能级')
        plt.ylabel('能量 E')
        plt.ylim(0, self.energy(max_n) * 1.1)
        plt.xticks([])  # 隐藏x轴刻度
        plt.grid(True, axis='y')
        plt.show(block=True)